package io.kiki.sba.registry.server.clientInterface.node.service;


import io.kiki.sba.registry.api.Channel;
import io.kiki.sba.registry.api.exchange.Callback;
import io.kiki.sba.registry.api.exchange.NodeExchanger;
import io.kiki.sba.registry.api.exchange.RequestException;
import io.kiki.sba.registry.api.exchange.message.SimpleRequest;
import io.kiki.sba.registry.common.model.ClientOffPublishers;
import io.kiki.sba.registry.common.model.ClientToClientInterfaceServerNodeId;
import io.kiki.sba.registry.common.model.Response;
import io.kiki.sba.registry.common.model.dataserver.*;
import io.kiki.sba.registry.common.model.slot.MultiSlotAccessGenericResponse;
import io.kiki.sba.registry.common.model.slot.Slot;
import io.kiki.sba.registry.common.model.slot.SlotAccessGenericResponse;
import io.kiki.sba.registry.common.model.store.DataCenterToSubDatumMap;
import io.kiki.sba.registry.common.model.store.Publisher;
import io.kiki.sba.registry.common.model.store.UnPublisher;
import io.kiki.sba.registry.common.model.store.Url;
import io.kiki.sba.registry.compress.CompressConstants;
import io.kiki.sba.registry.server.clientInterface.bootstrap.Config;
import io.kiki.sba.registry.server.clientInterface.slot.SlotsCache;
import io.kiki.sba.registry.server.shared.env.ServerEnv;
import io.kiki.sba.registry.server.shared.util.DatumUtils;
import io.kiki.sba.registry.task.BlockingQueues;
import io.kiki.sba.registry.task.FastRejectedExecutionException;
import io.kiki.sba.registry.task.MetricsableThreadPoolExecutor;
import io.kiki.sba.registry.task.RejectedDiscardHandler;
import io.kiki.sba.registry.util.ConcurrentUtils;
import io.kiki.sba.registry.util.OsUtils;
import io.kiki.sba.registry.util.StringFormatter;
import io.kiki.sba.registry.util.StringUtil;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;

import javax.annotation.PostConstruct;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;


public class DataStoreNodeServiceImpl implements DataStoreNodeService {

    private static final Logger logger = LoggerFactory.getLogger(DataStoreNodeServiceImpl.class);
    final RejectedDiscardHandler rejectedDiscardHandler = new RejectedDiscardHandler();
    private final ThreadPoolExecutor threadPoolExecutor = MetricsableThreadPoolExecutor.newExecutor("DataNodeCallback", OsUtils.getCpuCount() * 2, 4096, rejectedDiscardHandler);
    @Autowired
    private NodeExchanger dataStoreNodeExchanger;
    @Autowired
    private SlotsCache slotsCache;
    @Autowired
    private Config config;
    private Worker[] workers;
    private BlockingQueues<Request> queues;

    @PostConstruct
    public void init() {
        this.workers = new Worker[config.getDataNodeExecutorWorkerSize()];
        queues = new BlockingQueues<>(config.getDataNodeExecutorWorkerSize(), config.getDataNodeExecutorQueueSize(), false);
        for (int i = 0; i < workers.length; i++) {
            workers[i] = new Worker(queues.getQueue(i));
            ConcurrentUtils.createDaemonThread("req-data-worker-" + i, workers[i]).start();
        }
    }

    private void submit(int slotId, Request request) {
        int idx = slotId % queues.queueNum();
        try {
            queues.put(idx, request);
        } catch (FastRejectedExecutionException e) {
            throw new FastRejectedExecutionException(String.format("commit req overflow, slotId=%d, %s", slotId, e.getMessage()));
        }
    }

    @Override
    public void register(final Publisher publisher) {
        final int slotId = slotsCache.slotOf(publisher.getDataInfoId());
        submit(slotId, new Request(slotId, publisher));
    }

    @Override
    public void unregister(final Publisher publisher) {
        final int slotId = slotsCache.slotOf(publisher.getDataInfoId());
        UnPublisher unPublisher = UnPublisher.of(publisher);
        submit(slotId, new Request(slotId, unPublisher));
    }

    @Override
    public void clientOff(ClientOffPublishers clientOffPublishers) {
        if (clientOffPublishers.isEmpty()) {
            return;
        }
        Map<Integer, ClientOffPublisher> map = groupBySlot(clientOffPublishers);
        for (Map.Entry<Integer, ClientOffPublisher> entry : map.entrySet()) {
            final int slotId = entry.getKey();
            final ClientOffPublisher clientOffPublisher = entry.getValue();
            submit(slotId, new Request(slotId, clientOffPublisher));
        }
    }

    private Map<Integer, ClientOffPublisher> groupBySlot(ClientOffPublishers clientOffPublishers) {
        ClientToClientInterfaceServerNodeId clientToClientInterfaceServerNodeId = clientOffPublishers.getClientToClientInterfaceServerNodeId();
        List<Publisher> publisherList = clientOffPublishers.getPublishers();
        Map<Integer, ClientOffPublisher> map = new HashMap<>();
        for (Publisher publisher : publisherList) {
            final String dataInfoId = publisher.getDataInfoId();
            int slotId = slotsCache.slotOf(dataInfoId);
            ClientOffPublisher clientOffPublisher = map.computeIfAbsent(slotId, k -> new ClientOffPublisher(clientToClientInterfaceServerNodeId));
            clientOffPublisher.addPublisher(publisher);
        }
        return map;
    }

    @Override
    public void fetchDataVersion(String dataCenter, int slotId, Map<String, DatumVersion> dataInfoIdToDatumVersionMap, Callback<Map<String, DatumVersion>> callback) {
        final Slot slot = getSlot(dataCenter, slotId);
        final String leaderDataStoreNodeId = slot.getLeaderDataStoreNodeId();
        try {
            final GetDataVersionRequest request = new GetDataVersionRequest(dataCenter, ServerEnv.PROCESS_ID, slotId, dataInfoIdToDatumVersionMap);
            request.setSlotsEpoch(slotsCache.getEpoch(dataCenter));
            request.setSlotLeaderEpoch(slot.getLeaderEpoch());
            final io.kiki.sba.registry.api.Callback callback1 = new io.kiki.sba.registry.api.Callback() {
                @Override
                public void onSuccess(Channel channel, Object message) {
                    handleFetchDataVersionCallback(channel, message, slotId, leaderDataStoreNodeId, dataCenter, callback);
                }

                @Override
                public void onFailed(Channel channel, Throwable exception) {
                    callback.onFailed(channel, exception);
                }

                @Override
                public Executor getExecutor() {
                    return threadPoolExecutor;
                }
            };
            final Slot slot1 = getSlot(config.getDataCenter(), slotId);
            io.kiki.sba.registry.api.exchange.message.Request<GetDataVersionRequest> request1 = new SimpleRequest<>(getLeaderDataStoreNodeUrl(slot1), request, callback1);
            io.kiki.sba.registry.api.exchange.message.Response response = dataStoreNodeExchanger.request(request1);
            io.kiki.sba.registry.api.exchange.message.Response.ResultStatus resultStatus = (io.kiki.sba.registry.api.exchange.message.Response.ResultStatus) response.getResult();
            if (resultStatus != io.kiki.sba.registry.api.exchange.message.Response.ResultStatus.SUCCESSFUL) {
                throw new RequestException("response not success, status=" + resultStatus);
            }
        } catch (RequestException e) {
            throw new RuntimeException(StringFormatter.format("GetDataVersion fail {}@{}, slotId={}", leaderDataStoreNodeId, dataCenter, slotId, e));
        }
    }

    void handleFetchDataVersionCallback(Channel channel, Object message, int slotId, String dataNodeIp, String dataCenter, Callback<Map<String, DatumVersion>> callback) {
        SlotAccessGenericResponse<Map<String, DatumVersion>> slotAccessGenericResponse = (SlotAccessGenericResponse<Map<String, DatumVersion>>) message;
        if (slotAccessGenericResponse.isSuccess()) {
            Map<String, DatumVersion> dataInfoIdToDatumVersionMap = slotAccessGenericResponse.getData();
            DatumUtils.intern(dataInfoIdToDatumVersionMap);
            callback.onSuccess(channel, dataInfoIdToDatumVersionMap);
        } else {
            callback.onFailed(channel, new RuntimeException(StringFormatter.format("GetDataVersion failed, {}@{}, slotId={}, access={}, msg:{}", dataNodeIp, dataCenter, slotId, slotAccessGenericResponse.getSlotAccess(), slotAccessGenericResponse.getMessage())));
        }
    }

    @Override
    public DataCenterToSubDatumMap fetch(Set<String> dataCenters, String dataInfoId) {
        final Slot localSlot = getSlot(config.getDataCenter(), dataInfoId);
        int slotId = localSlot.getId();
        String leaderDataStoreNodeId = localSlot.getLeaderDataStoreNodeId();

        Map<String, Long> slotsEpochs = new HashMap<>(dataCenters.size());
        Map<String, Long> slotLeaderEpochs = new HashMap<>(dataCenters.size());
        for (String dataCenter : dataCenters) {
            final Slot slot = getSlot(dataCenter, dataInfoId);
            StringUtil.checkEquals(slotId, slot.getId(), "slotId");
            slotsEpochs.put(dataCenter, slotsCache.getEpoch(dataCenter));
            slotLeaderEpochs.put(dataCenter, slot.getLeaderEpoch());
        }
        try {
            GetMultiDataRequest request1 = new GetMultiDataRequest(ServerEnv.PROCESS_ID, slotId, dataInfoId, CompressConstants.defaultCompressEncodes, slotsEpochs, slotLeaderEpochs);

            io.kiki.sba.registry.api.exchange.message.Request<GetMultiDataRequest> request = new io.kiki.sba.registry.api.exchange.message.Request<GetMultiDataRequest>() {

                @Override
                public GetMultiDataRequest getRequestBody() {
                    return request1;
                }

                @Override
                public Url getRequestUrl() {
                    return getLeaderDataStoreNodeUrl(localSlot);
                }

                @Override
                public Integer getTimeout() {
                    return config.getDataNodeExchangeForFetchDatumTimeoutMillis();
                }
            };

            io.kiki.sba.registry.api.exchange.message.Response response = dataStoreNodeExchanger.request(request);
            Object result = response.getResult();
            MultiSlotAccessGenericResponse<DataCenterToSubDatumMap> genericResponse = (MultiSlotAccessGenericResponse<DataCenterToSubDatumMap>) result;
            if (genericResponse.isSuccess()) {
                final DataCenterToSubDatumMap dataCenterToSubDatumMap = genericResponse.getData();
                if (dataCenterToSubDatumMap == null) {
                    return null;
                }
                return DataCenterToSubDatumMap.intern(dataCenterToSubDatumMap);
            } else {
                throw new RuntimeException(StringFormatter.format("GetMultiData got fail response {}, {}, {}, slotId={} msg:{}", leaderDataStoreNodeId, dataInfoId, dataCenters, slotId, genericResponse.getMessage()));
            }
        } catch (RequestException e) {
            throw new RuntimeException(StringFormatter.format("GetMultiData fail {}, {}, {}, slotId={}", leaderDataStoreNodeId, dataInfoId, dataCenters, slotId), e);
        }
    }

    private Response sendRequest(io.kiki.sba.registry.api.exchange.message.Request request) throws RequestException {
        io.kiki.sba.registry.api.exchange.message.Response response = dataStoreNodeExchanger.request(request);
        Object result = response.getResult();
        SlotAccessGenericResponse resp = (SlotAccessGenericResponse) result;
        if (!resp.isSuccess()) {
            throw new RuntimeException(String.format("response failed, target: %s, request: %s, message: %s", request.getRequestUrl(), request.getRequestBody(), resp.getMessage()));
        }
        return resp;
    }

    private Slot getSlot(String dataCenter, String dataInfoId) {
        Slot slot = slotsCache.getSlot(dataCenter, dataInfoId);
        if (slot == null) {
            throw new RequestException(StringFormatter.format("slot not found for dataCenter={}, dataInfoId={}", dataCenter, dataInfoId));
        }
        return slot;
    }

    private Slot getSlot(String dataCenter, int slotId) {
        Slot slot = slotsCache.getSlot(dataCenter, slotId);
        if (slot == null) {
            throw new RequestException(StringFormatter.format("slot not found, slotId={}", slotId));
        }
        return slot;
    }

    private Url getLeaderDataStoreNodeUrl(Slot slot) {
        final String leaderDataStoreNodeId = slot.getLeaderDataStoreNodeId();
        if (StringUtils.isBlank(leaderDataStoreNodeId)) {
            throw new RequestException(String.format("slot has no leader, slotId=%s", slot));
        }
        return new Url(leaderDataStoreNodeId, config.getDataServerPort());
    }


    private boolean request(BatchRequest batchRequest) {
        try {
            String localDataCenter = config.getDataCenter();
            final Slot slot = getSlot(localDataCenter, batchRequest.getSlotId());
            batchRequest.setSlotsEpoch(slotsCache.getEpoch(localDataCenter));
            batchRequest.setSlotLeaderEpoch(slot.getLeaderEpoch());
            sendRequest(new io.kiki.sba.registry.api.exchange.message.Request() {
                @Override
                public Object getRequestBody() {
                    return batchRequest;
                }

                @Override
                public Url getRequestUrl() {
                    return getLeaderDataStoreNodeUrl(slot);
                }
            });
            return true;
        } catch (Throwable e) {
            logger.error("failed to request batch, {}", batchRequest, e);
            return false;
        }
    }

    private Map<Integer, LinkedList<Object>> drainTo(BlockingQueue<Request> queue, int max) {
        List<Request> list = new ArrayList<>(max);
        queue.drainTo(list, max);
        if (list.isEmpty()) {
            return Collections.emptyMap();
        }
        Map<Integer, LinkedList<Object>> map = new LinkedHashMap<>(list.size());
        for (Request request : list) {
            LinkedList<Object> objects = map.computeIfAbsent(request.slotId, k -> new LinkedList<>());
            objects.add(request.request);
        }
        return map;
    }

    private static final class Request {
        final int slotId;
        final Object request;

        Request(int slotId, Object request) {
            this.slotId = slotId;
            this.request = request;
        }
    }

    private static final class RetryBatchRequest {
        final BatchRequest batchRequest;
        long expireTimestamp;
        int retryCount;

        RetryBatchRequest(BatchRequest batchRequest) {
            this.batchRequest = batchRequest;
        }


    }

    private final class Worker implements Runnable {
        final BlockingQueue<Request> queue;
        final LinkedList<RetryBatchRequest> retryBatchRequestList = new LinkedList<>();

        Worker(BlockingQueue<Request> queue) {
            this.queue = queue;
        }

        @Override
        public void run() {
            for (; ; ) {
                try {
                    final Request firstRequest = queue.poll(200, TimeUnit.MILLISECONDS);
                    if (firstRequest != null) {
                        // TODO config max
                        Map<Integer, LinkedList<Object>> map = drainTo(queue, config.getDataNodeMaxBatchSize());
                        // send by order, firstReq.slotId is the first one
                        LinkedList<Object> list = map.remove(firstRequest.slotId);
                        if (list == null) {
                            list = new LinkedList<>();
                        }
                        list.addFirst(firstRequest.request);
                        request(firstRequest.slotId, list);
                        //
                        for (Map.Entry<Integer, LinkedList<Object>> entry : map.entrySet()) {
                            request(entry.getKey(), entry.getValue());
                        }
                    }
                    // check the retry
                    if (!retryBatchRequestList.isEmpty()) {
                        final Iterator<RetryBatchRequest> iterator = retryBatchRequestList.iterator();
                        List<RetryBatchRequest> list = new ArrayList<>();
                        while (iterator.hasNext()) {
                            RetryBatchRequest retryBatchRequest = iterator.next();
                            iterator.remove();
                            if (!DataStoreNodeServiceImpl.this.request(retryBatchRequest.batchRequest)) {
                                list.add(retryBatchRequest);
                            }
                        }
                        for (RetryBatchRequest retryBatchRequest : list) {
                            retry(retryBatchRequest);
                        }
                    }
                } catch (Throwable e) {
                    logger.error("failed to request batch", e);
                }
            }
        }

        private boolean request(int slotId, List<Object> list) {
            final BatchRequest batchRequest = new BatchRequest(ServerEnv.PROCESS_ID, slotId, list);
            if (!DataStoreNodeServiceImpl.this.request(batchRequest)) {
                retry(new RetryBatchRequest(batchRequest));
                return false;
            }
            return true;
        }

        private boolean retry(RetryBatchRequest retryBatchRequest) {
            retryBatchRequest.retryCount++;
            if (retryBatchRequest.retryCount <= config.getDataNodeRetryTimes()) {
                if (retryBatchRequestList.size() >= config.getDataNodeRetryQueueSize()) {
                    // remove the oldest
                    retryBatchRequestList.removeFirst();
                }
                retryBatchRequest.expireTimestamp = System.currentTimeMillis() + config.getDataNodeRetryBackoffMillis();
                retryBatchRequestList.add(retryBatchRequest);
                return true;
            }
            return false;
        }


    }
}
